# -*- coding:utf-8 -*-


'''
使用tnsorflow全连接神经网络
'''

# -*- coding:utf-8 -*-

import tensorflow as tf
import library.img_data as idata
from numpy import *
import os
import time


# 数据集类
pre = os.getcwd() + os.path.sep + ".." + os.path.sep
vertifyImgs = idata.VertifyImg(pre + "trainimg",pre +  "testimg")


# 获取一个图片向量的分类结果
def getRes1(inX, W1, b1, W2, b2):
    sess = tf.InteractiveSession()
    inX = mat(inX)
    x = tf.convert_to_tensor(inX)
    W1 = tf.convert_to_tensor(W1)
    b1 =  tf.convert_to_tensor(b1)
    W2 =  tf.convert_to_tensor(W2)
    b2 =  tf.convert_to_tensor(b2)
    hidden1 = tf.nn.relu(tf.matmul(x, W1) + b1)
    y = tf.nn.softmax(tf.matmul(hidden1, W2) + b2)
    tf.global_variables_initializer().run()
    index = tf.argmax(y, 1).eval()[0]
    labelIndexDir = idata.getLabelIndexDir()
    i = list(labelIndexDir.keys())[list(labelIndexDir.values()).index(index)]
    sess.close()
    return i

# 获取一个图片路径所表示的图片的分类结果
def getRes2(imgPath, W1, b1, W2, b2):
    inX = mat(idata.img2Vec(imgPath))
    return getRes1(inX, W1, b1, W2, b2)

# 在训练数据集上进行训练
def startTrain(kp = 0.75):
    sess = tf.InteractiveSession()
    in_units = 300
    h1_units = 100
    out_units = 33
    W1 = tf.Variable(tf.truncated_normal([in_units, h1_units], stddev=0.1))
    b1 = tf.Variable(tf.zeros([h1_units]))
    W2 = tf.Variable(tf.zeros([h1_units, out_units]))
    b2 = tf.Variable(tf.zeros([out_units]))
    x = tf.placeholder(tf.float32, [None, in_units])
    keep_prob = tf.placeholder(tf.float32)
    hidden1 = tf.nn.relu(tf.matmul(x, W1) + b1)
    hidden1_drop = tf.nn.dropout(hidden1, keep_prob)
    y = tf.nn.softmax(tf.matmul(hidden1_drop, W2) + b2)
    y_ = tf.placeholder(tf.float32, [None, out_units])
    cross_entropy = tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(y), reduction_indices=[1]))
    train_step = tf.train.AdagradOptimizer(0.1).minimize(cross_entropy)
    tf.global_variables_initializer().run()
    for i in range(1000):
        batch_xs, batch_ys = vertifyImgs.next_batch(100)
        train_step.run({x: batch_xs, y_: batch_ys, keep_prob: kp})
        if i % 100 == 0:
            cross = cross_entropy.eval({x: batch_xs, y_: batch_ys, keep_prob: 1.0})
            print("repeat: %d " % i, "cross: %.4f" % cross)
    W1, b1, W2, b2 = sess.run((W1, b1, W2, b2))
    sess.close()
    return  W1, b1, W2, b2

# 在测试数据集上进行测试
def startTest(W1, b1, W2, b2):
    sess = tf.InteractiveSession()
    testDataSet, testLables = vertifyImgs.get_test_data()
    x = tf.convert_to_tensor(testDataSet)
    W1 = tf.convert_to_tensor(W1)
    b1 = tf.convert_to_tensor(b1)
    W2 = tf.convert_to_tensor(W2)
    b2 = tf.convert_to_tensor(b2)
    hidden1 = tf.nn.relu(tf.matmul(x, W1) + b1)
    y = tf.nn.softmax(tf.matmul(hidden1, W2) + b2)
    y_ = tf.convert_to_tensor(testLables)
    correct_prediction = tf.equal(tf.argmax(y, 1), tf.argmax(y_, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    tf.global_variables_initializer().run()
    print("correct rate:%.2f" % (accuracy.eval()))
    sess.close()

# 获取一张原始图片的识别结果
def vertifyInitimg(imgPath, W1, b1, W2, b2):
    sonFilePaths = idata.getAllSonImg(imgPath)
    res = ''
    for sonFilePath in sonFilePaths:
        res = res + getRes2(sonFilePath, W1, b1, W2, b2)
    return res



if __name__ == '__main__':
    W1, b1, W2, b2 = startTrain()
    res = vertifyInitimg('initimg' +  os.path.sep +'0.png', W1, b1, W2, b2)
    print(res)
